'''
第一关基础测试
'''
W0 = ""
W1 = ""
W2 = ""
W3 = ""
W4 = ""
W5 = ""
S1_box = [["1001", "0100", "1010", "1011"], ["1101", "0001", "1000", "0101"],
           ["0110", "0010", "0000", "0011"], ["1100", "1110", "1111", "0111"]]  # S盒
S2_box = [["1010", "0101", "1001", "1011"], ["0001", "0111", "1000", "1111"],
           ["0110", "0000", "0010", "0011"], ["1100", "0100", "1101", "1110"]]  # 逆S盒

def get_Key():
    '''
    输入主密匙，获取K1,K2
    '''
    print("请输入主密匙(16bit):")
    main_Key = input()
    global W0
    global W1
    global W2
    global W3
    global W4
    global W5
    W0 = main_Key[0:8]
    W1 = main_Key[8:16]
    W2 = xor(W0, g(W1, 1))
    W3 = xor(W1, W2)
    W4 = xor(W2, g(W3, 2))
    W5 = xor(W4, W3)
    print("w2：" + W2)
    print("w3：" + W3)
    print("w4：" + W4)
    print("w5：" + W5)

def encrypt():
    '''
    加密函数
    '''
    print("请输入加密信息(16bit):")
    Plainttext = input() #明文
    #第一轮
    Plainttext = xor(Plainttext, W0+W1) #轮秘钥加
    Plainttext = replace(Plainttext, 1) #半字节替代
    Plainttext = move_hang(Plainttext) #行移位
    Plainttext = multiply_lie(Plainttext, 1) #列混淆
    #第二轮
    Plainttext = xor(Plainttext, W2+W3) #轮秘钥加
    Plainttext = replace(Plainttext, 1) #半字节替代
    Plainttext = move_hang(Plainttext) #行移位
    ciphertext = xor(Plainttext, W4+W5) #轮秘钥加
    print("密文为: " + ciphertext)

def decrypt():
    '''
    解密函数
    '''
    print("请输入解密信息(16bit):")
    ciphertext = input() #秘文
    # 轮秘钥加
    ciphertext = xor(ciphertext, W4+W5)
    # 行移位
    ciphertext = move_hang(ciphertext)
    # 半字节替代
    ciphertext = replace(ciphertext, 2)
    # 轮秘钥加
    ciphertext = xor(ciphertext, W2+W3)
    # 列混淆
    ciphertext = multiply_lie(ciphertext, 2)
    # 行移位
    ciphertext = move_hang(ciphertext)
    # 半字节替代
    ciphertext = replace(ciphertext, 2)
    Plainttext = xor(ciphertext, W0+W1) # 轮秘钥加

    print("明文为: " + Plainttext)

def g(key, n):
    '''
    根据key生成轮密钥
    n代表的是第几轮
    '''
    N0 = key[0:4]
    N1 = key[4:8]
    N1_ = searchSbox(N1, 1)
    N0_ = searchSbox(N0, 1)
    N_ = N1_ + N0_
    W_ = ""
    if n == 1:
        W_ = xor(N_, "10000000")
    elif n == 2:
        W_ = xor(N_, "00110000")
    return W_

def searchSbox(str, n):
    '''
    用于在 S1_box 或 S2_box 中查找对应的值
    输入:str,n
    它接收一个二进制字符串  str  和一个整数  n  作为参数，根据  n  的值选择要查找的 S_box，并返回对应的值。
    '''
    ret = str[0] + str[1]
    ret1 = str[2] + str[3]
    retu = ""
    if n == 1:
        retu = S1_box[int(ret, 2)][int(ret1, 2)]
    else:
        retu = S2_box[int(ret, 2)][int(ret1, 2)]
    return retu

def xor(str, key):
    '''
    异或运算
    '''
    sb = []
    for i in range(len(str)):
        if str[i] == key[i]:
            sb.append("0")
        else:
            sb.append("1")
    return "".join(sb)


def replace(str, n):
    '''
    半字节替换
    n代表半字节替换或是逆替换
    '''
    str0 = str[0:4]
    str1 = str[4:8]
    str2 = str[8:12]
    str3 = str[12:16]
    str0 = searchSbox(str0, n)
    str1 = searchSbox(str1, n)
    str2 = searchSbox(str2, n)
    str3 = searchSbox(str3, n)
    return str0 + str1 + str2 + str3

def move_hang(str):
    '''
    行移位
    '''
    str0=str[0:4]
    str1=str[4:8]
    str2=str[8:12]
    str3=str[12:16]
    str_=str1
    str1=str3
    str3=str_
    return str0+str1+str2+str3

def multiply_lie(str, n):
    '''
    列混淆
    '''
    s0 = str[0:4]
    s1 = str[4:8]
    s2 = str[8:12]
    s3 = str[12:16]
    s0_ = "0000"
    s1_ = "0000"
    s2_ = "0000"
    s3_ = "0000"
    if n == 1:
        s1_4 = multiply(s1, 4)
        s0_ = xor(s0, multiply(s1, 4))
        s1_ = xor(s1, multiply(s0, 4))
        s2_ = xor(s2, multiply(s3, 4))
        s3_ = xor(s3, multiply(s2, 4))
    elif n == 2:
        s0_ = xor(multiply(s0, 9), multiply(s1, 2))
        s1_ = xor(multiply(s0, 2), multiply(s1, 9))
        s2_ = xor(multiply(s2, 9), multiply(s3, 2))
        s3_ = xor(multiply(s2, 2), multiply(s3, 9))
    return s0_ + s1_ + s2_ + s3_

def multiply(str, n):
    '''
    乘法
    '''
    result = ""
    if n == 2:
        if str == "0000":
            result = "0000"
        elif str == "0001":
            result = "0010"
        elif str == "0010":
            result = "0100"
        elif str == "0011":
            result = "0110"
        elif str == "0100":
            result = "1000"
        elif str == "0101":
            result = "1010"
        elif str == "0110":
            result = "1100"
        elif str == "0111":
            result = "1110"
        elif str == "1000":
            result = "0011"
        elif str == "1001":
            result = "0001"
        elif str == "1010":
            result = "0111"
        elif str == "1011":
            result = "0101"
        elif str == "1100":
            result = "1011"
        elif str == "1101":
            result = "1001"
        elif str == "1110":
            result = "1111"
        elif str == "1111":
            result = "1101"
        
    if n == 4:
        if str == "0000":
            result = "0000"
        elif str == "0001":
            result = "0100"
        elif str == "0010":
            result = "1000"
        elif str == "0011":
            result = "1100"
        elif str == "0100":
            result = "0011"
        elif str == "0101":
            result = "0111"
        elif str == "0110":
            result = "1011"
        elif str == "0111":
            result = "1111"
        elif str == "1000":
            result = "0110"
        elif str == "1001":
            result = "0010"
        elif str == "1010":
            result = "1110"
        elif str == "1011":
            result = "1010"
        elif str == "1100":
            result = "0101"
        elif str == "1101":
            result = "0001"
        elif str == "1110":
            result = "1101"
        elif str == "1111":
            result = "1001"
        
    if n == 9:
        if str == "0000":
            result = "0000"
        elif str == "0001":
            result = "1001"
        elif str == "0010":
            result = "0001"
        elif str == "0011":
            result = "1000"
        elif str == "0100":
            result = "0010"
        elif str == "0101":
            result = "1011"
        elif str == "0110":
            result = "0011"
        elif str == "0111":
            result = "1010"
        elif str == "1000":
            result = "0100"
        elif str == "1001":
            result = "1101"
        elif str == "1010":
            result = "0101"
        elif str == "1011":
            result = "1100"
        elif str == "1100":
            result = "0110"
        elif str == "1101":
            result = "1110"
        elif str == "1110":
            result = "0111"
        elif str == "1111":
            result = "1111"
        
    return result

if __name__ == "__main__": 
    get_Key()
    encrypt()
    decrypt()
    '''
    测试输出
    请输入主密匙(16bit):
    1111000011110000
    w2：11100111
    w3：00010111
    w4：10000011
    w5：10010100
    请输入加密信息(16bit):
    1111111111111111
    密文为: 1110100111001110
    请输入解密信息(16bit):
    1110100111001110
    明文为: 1111111111111111
    '''
